package net.anquankong.websec.argus;

import java.io.*;
import java.net.*;
import java.util.*;
import java.security.*;

import com.gargoylesoftware.htmlunit.*;
import com.gargoylesoftware.htmlunit.html.*;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.commons.io.*;
import org.apache.commons.codec.binary.Base64;

public class WebClientThread extends Thread {
	// Private vars
	// private ArgusContext argusContext_ = null;
	private String uniqueId_;
	private String url_;
	private Boolean enableJavaScript_;
	private Boolean enableFlashParser_;
	private WebClient webClient_ = null;
	private Object pageLoadFinishedLock_ = new Object();
	private Boolean pageLoadFinished_ = false;
	private String tmpFolder_;
	private MessageDigest sha1Digest;

	private static Log log_ = LogFactory.getLog(WebClientThread.class);

	// Sub-classes to handle prompt and confirmation box
	public static class DefaultPromptHandler extends Object implements
			PromptHandler {
		public String handlePrompt(Page page, String message) {
			return "";
		}
	}

	public static class DefaultConfirmHandler extends Object implements
			ConfirmHandler {
		public boolean handleConfirm(Page page, String message) {
			return true;
		}
	}

	WebClientThread(/* ArgusContext argusContext, */String url,
			Boolean enableJavaScript, Boolean enableFlashParser,
			String tmpFolder) {
		// argusContext_ = argusContext;
		url_ = url;
		enableJavaScript_ = enableJavaScript;
		enableFlashParser_ = enableFlashParser;
		tmpFolder_ = tmpFolder;

		try {
			sha1Digest = MessageDigest.getInstance("SHA1");
		} catch (NoSuchAlgorithmException e) {
			// TODO Auto-generated catch block
		}
	}

	public String serializeResults(Boolean outputAnchors, Boolean outputAST,
			Boolean outputText, Boolean outputObjects) {
		StringBuilder resultBuilder = new StringBuilder();

		if (outputAnchors) {
			synchronized (webClient_.argusAnchorsListLock) {
				resultBuilder.append(String.format("Anchors %d%n",
						webClient_.argusAnchorList.size()));
				for (String anchorStr : webClient_.argusAnchorList) {
					anchorStr = anchorStr.replace("\r", "");
					anchorStr = anchorStr.replace("\n", "");
					resultBuilder.append(String.format("%s%n", anchorStr));
				}
			}
		}

		if (outputAST) {
			synchronized (webClient_.argusASTListLock) {
				StringBuilder astBuilder = new StringBuilder();
				for (String astStr : webClient_.argusASTList) {
					astBuilder.append(astStr);
				}

				String ast = astBuilder.toString();
				int astLength;
				try {
					astLength = ast.length();
				} catch (Exception e) {
					astLength = 0;
					ast = "";
				}
				
				// BASE64
				String astBase64Str = "";
				try {
					byte[] astBytes = ast.getBytes("UTF8");
					byte[] astBase64 = Base64.encodeBase64(astBytes);
					astBase64Str = new String(astBase64, "UTF8");
				} catch (UnsupportedEncodingException e) {
					// TODO Auto-generated catch block					
				}
				

				resultBuilder.append(String.format("AST %d%n", astLength));
				resultBuilder.append(String.format("%s%n", astBase64Str));
			}
		}

		if (outputText) {
			synchronized (webClient_.argusTextListLock) {
				StringBuilder textBuilder = new StringBuilder();
				for (String textStr : webClient_.argusTextList) {
					textBuilder.append(String.format("%s%n", textStr));
				}

				String text = textBuilder.toString();
				int textLength;
				try {
					textLength = text.length();
				} catch (Exception e) {
					textLength = 0;
					text = "";
				}

				// BASE64
				String textBase64Str = "";
				try {
					byte[] textBytes = text.getBytes("UTF8");
					byte[] textBase64 = Base64.encodeBase64(textBytes);
					textBase64Str = new String(textBase64, "UTF8");
				} catch (UnsupportedEncodingException e) {
					// TODO Auto-generated catch block					
				}

				resultBuilder.append(String.format("Text %d%n", textLength));
				resultBuilder.append(String.format("%s%n", textBase64Str));
			}
		}

		if (outputObjects) {
			synchronized (webClient_.argusObjectListLock) {
				resultBuilder.append(String.format("Objects %d%n",
						webClient_.argusObjectList.size()));
				for (String objectStr : webClient_.argusObjectList) {
					resultBuilder.append(String.format("%s%n", objectStr));
				}
			}

		}

		return resultBuilder.toString();
	}

	public String serializeResults() {
		StringBuilder resultBuilder = new StringBuilder();

		synchronized (webClient_.argusAnchorsListLock) {
			resultBuilder.append(String.format("Anchors %d%n",
					webClient_.argusAnchorList.size()));
			for (String anchorStr : webClient_.argusAnchorList) {
				resultBuilder.append(String.format("%s%n", anchorStr));
			}
		}

		synchronized (webClient_.argusASTListLock) {
			StringBuilder astBuilder = new StringBuilder();
			for (String astStr : webClient_.argusASTList) {
				astBuilder.append(astStr);
			}

			String ast = astBuilder.toString();
			byte[] astBytes;
			int astLength;
			try {
				astLength = ast.length();
			} catch (Exception e) {
				astLength = 0;
				ast = "";
			}

			resultBuilder.append(String.format("AST %d%n", astLength));
			resultBuilder.append(ast);
		}

		synchronized (webClient_.argusTextListLock) {
			StringBuilder textBuilder = new StringBuilder();
			for (String textStr : webClient_.argusTextList) {
				textBuilder.append(String.format("%s%n", textStr));
			}

			String text = textBuilder.toString();
			byte[] textBytes;
			int textLength;
			try {
				textLength = text.length();
			} catch (Exception e) {
				textLength = 0;
				text = "";
			}

			resultBuilder.append(String.format("Text %d%n", textLength));
			resultBuilder.append(text);
		}

		return resultBuilder.toString();
	}

	public boolean isFinished() {
		return pageLoadFinished_;
	}

	public void interrupt() {
		synchronized (pageLoadFinishedLock_) {
			if (!pageLoadFinished_) {
				try {
					webClient_.closeAllWindows();
					super.interrupt();
				} catch (Exception e) {
					// TODO Auto-generated catch block
				}
			}
		}
	}

	public void run() {
		List<String> collectedAlerts = new ArrayList<String>();

		webClient_ = new WebClient(BrowserVersion.INTERNET_EXPLORER_8);
		webClient_.setAlertHandler(new CollectingAlertHandler(collectedAlerts));
		webClient_.setPromptHandler(new DefaultPromptHandler());
		webClient_.setConfirmHandler(new DefaultConfirmHandler());
		webClient_.setThrowExceptionOnScriptError(false);
		webClient_.setThrowExceptionOnFailingStatusCode(false);
		webClient_.setAjaxController(new NicelyResynchronizingAjaxController());

		// Enable/disable script
		webClient_.setJavaScriptEnabled(enableJavaScript_);

		log_.error("[PARSER] " + uniqueId_ + " started.");
		try {
			HtmlPage page = webClient_.getPage(url_);

			// Anchor output
			// Make temp folder if we need Flash parser
			if (enableFlashParser_) {
				(new File(tmpFolder_)).mkdir();
			}
			HashSet<String> anchorUrlList = extractAnchorFromHtmlPage(page, 0,
					2);
			for (String anchorUrl : anchorUrlList) {
				webClient_.argusAnchorList.add(anchorUrl);
			}

			// Text output
			String pageText = page.asText();
			webClient_.argusTextList.add(pageText);
			String pageHtml = page.asXml();
			webClient_.argusTextList.add(pageHtml);
		} catch (Exception e) {
			// TODO Auto-generated catch block
		} finally {
			try {
				webClient_.closeAllWindows();
			} catch (Exception e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}

			try {
				FileUtils.deleteDirectory(new File(tmpFolder_));
			} catch (Exception e) {
				// TODO Auto-generated catch block
			}

			log_.error("[PARSER] " + uniqueId_ + " closed.");
		}

		synchronized (pageLoadFinishedLock_) {
			pageLoadFinished_ = true;
		}

		// Send finishing signal
		// ClientUDPSender.sendCTRLCommand(this.argusContext_.dest,
		// this.argusContext_.port,
		// ClientUDPSender.ClientUDPSenderCTRLCommandTypes.FINISHED,
		// this.argusContext_.pipeIdx);
	}

	public static String getHex(byte[] raw) {
		String HEXES = "0123456789abcdef";
		if (raw == null) {
			return null;
		}
		final StringBuilder hex = new StringBuilder(2 * raw.length);
		for (final byte b : raw) {
			hex.append(HEXES.charAt((b & 0xF0) >> 4)).append(
					HEXES.charAt((b & 0x0F)));
		}
		return hex.toString();
	}

	private HashSet<String> extractAnchorFromHtmlPage(HtmlPage page, int level,
			int levelLimit) {
		if (level > levelLimit)
			return null; // Exceeds frame recursive level

		HashSet<String> urlList = new HashSet<String>();

		URL baseUrl = page.getUrl();

		List<HtmlAnchor> anchorList = page.getAnchors();
		for (HtmlAnchor anchorElement : anchorList) {
			String hrefAttr = anchorElement.getHrefAttribute();
			try {
				if (!hrefAttr.isEmpty()) {
					URL anchorUrl = new URL(baseUrl, hrefAttr);
					urlList.add(anchorUrl.toString());
				}
			} catch (Exception ex) {
			}
		}

		// Enum all frames and get its page object
		List<FrameWindow> frameWinList = page.getFrames();
		for (FrameWindow frameWin : frameWinList) {
			HashSet<String> frameUrlList = extractAnchorFromHtmlPage(
					frameWin.getEnclosingPage(), level + 1, levelLimit);

			if (frameUrlList != null) {
				for (String frameUrl : frameUrlList) {
					urlList.add(frameUrl);
				}
			}
		}

		// Enum Flash objects
		if (enableFlashParser_) {
			byte[] readBuf = new byte[153600];

			DomNodeList<HtmlElement> embedElements = page
					.getElementsByTagName("embed");
			for (HtmlElement embedElement : embedElements) {
				String typeAttr = embedElement.getAttribute("type");
				if (typeAttr.indexOf("x-shockwave-flash") != -1) // "application/x-shockwave-flash"
				{
					String srcAttr = embedElement.getAttribute("src");
					try {
						URL SwfUrl = new URL(baseUrl, srcAttr);

						// Calc sha-1 hash of this url to name the output file
						sha1Digest.update(SwfUrl.toString().getBytes());
						byte[] urlHashByteArray = sha1Digest.digest();
						String urlHashStr = getHex(urlHashByteArray);

						String outputFilename = tmpFolder_ + urlHashStr
								+ ".swf";
						File fTest = new File(outputFilename);
						if (!fTest.exists()) {

							HttpURLConnection SwfConn = (HttpURLConnection) SwfUrl
									.openConnection();
							SwfConn.setReadTimeout(15000); // set read time-out
															// to 15 secs

							int httpStatus = SwfConn.getResponseCode();
							if (httpStatus == HttpURLConnection.HTTP_OK) {
								InputStream SwfStream = SwfConn
										.getInputStream();

								FileOutputStream writer = new FileOutputStream(
										outputFilename);
								int bytesRead;
								int totalBytesRead = 0;

								while ((bytesRead = SwfStream.read(readBuf)) > 0) {
									writer.write(readBuf, 0, bytesRead);
									totalBytesRead += bytesRead;

									if (totalBytesRead > readBuf.length * 10)
										break;
								}

								// System.err.println(totalBytesRead
								// + " bytes received.");

								writer.close();
								SwfStream.close();

								if (totalBytesRead <= readBuf.length * 10) {
									List<String> command = new ArrayList<String>();
									command.add("swfdump.exe");
									command.add("-a");
									command.add(outputFilename);

									ProcessBuilder builder = new ProcessBuilder(
											command);
									Map<String, String> environ = builder
											.environment();
									builder.directory(new File("."));

									final Process process = builder.start();
									InputStream is = process.getInputStream();
									InputStreamReader isr = new InputStreamReader(
											is, "UTF8");
									BufferedReader br = new BufferedReader(isr);
									String line;
									while ((line = br.readLine()) != null) {
										int i = line.indexOf("GetUrl URL:");
										if (i != -1) {
											i = i + ("GetUrl URL:").length();
											line = line.substring(i);
											String[] parts = line.split("\"");
											String relativeUrl = parts[1];

											try {
												// URL anchorUrl = new
												// URL(SwfUrl,
												// relativeUrl);
												URL anchorUrl = new URL(
														baseUrl, relativeUrl);
												urlList.add(anchorUrl
														.toString());
											} catch (Exception ex) {

											}
										}
									}
									process.waitFor();
								}
							}
						} else {
							// System.err.println("Flash object url: "
							// + SwfUrl.toString() + " in page: "
							// + baseUrl.toString() + " already exists.");
						}
					} catch (Exception ex) {
					}
				}
			}
		}

		return urlList;
	}
}
